package t20240514_lineage;

import com.alibaba.fastjson.JSONObject;
import io.debezium.util.Strings;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.metadata.RelColumnOrigin;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexTableInputRef;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.api.bridge.java.internal.StreamTableEnvironmentImpl;

import org.apache.flink.table.operations.CatalogSinkModifyOperation;
import org.apache.flink.table.operations.Operation;
import org.apache.flink.table.operations.ddl.CreateTableOperation;
import org.apache.flink.table.planner.operations.PlannerQueryOperation;
import org.apache.flink.table.planner.plan.schema.TableSourceTable;

import java.util.*;

public class Test1 {

    public static final String SQL_EMPTY_STR = "[\\s\\t\\n\\r]";


    public static void main(String[] args) {



        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        StreamTableEnvironment tbEnv = StreamTableEnvironment.create(env);


        String sourceSql = "CREATE TABLE datagen (" +
                "    score INT," +
                "    sn STRING" +
                "  ) WITH (" +
                "    'connector' = 'datagen'," +
                "    'rows-per-second'='1'," +
                "    'fields.score.kind'='sequence'," +
                "    'fields.score.start'='1'," +
                "    'fields.score.end'='10'," +
                "    'fields.sn.length'='10'" +
                "  )";

        String sinkSql =
                "  CREATE TABLE redis_sink (" +
                        "    key STRING," +
                        "    score INT," +
                        "    sn STRING," +
                        "    mins DOUBLE," +
                        "    maxs DOUBLE" +
                        "    ) WITH (" +
                        "    'connector' = 'redis'," +
                        "    'host' = '127.0.0.1'," +
                        "    'port' = '6379'," +
                        "    'redis-mode' = 'single'," +
                        "    'timeout' = '2000'," +
                        "    'password' = 'yy2024'," +
                        "    'command' = 'ZADD'," +
                        "    'zrem.rangeby' = 'SCORE'," +
                        "    'max.retries' = '5'" +
                        "  )";

        String execSql = "INSERT INTO redis_sink select 'SN_TEST' as key,score,sn,1,5 from datagen";



        parseSql(tbEnv, sourceSql);
        parseSql(tbEnv, sinkSql);
        parseSql(tbEnv, execSql);


    }

    private static void parseSql(StreamTableEnvironment tbEnv, String singleSql) {
        SqlType operationType = getOperationType(singleSql);
        StreamTableEnvironmentImpl tbEnvImpl = (StreamTableEnvironmentImpl) tbEnv;
        if (operationType.equals(SqlType.INSERT)) {


            List<Operation> list = tbEnvImpl.getParser().parse(singleSql);

            Operation operation = list.get(0);
            System.out.println(operation);



            CatalogSinkModifyOperation sinkOperation = (CatalogSinkModifyOperation) operation;

            PlannerQueryOperation queryOperation = (PlannerQueryOperation) sinkOperation.getChild();
            RelNode relNode = queryOperation.getCalciteTree();

            String sinkTbName = sinkOperation.getTableIdentifier().asSummaryString();
            System.out.println("sinkTbName:" + sinkTbName);
            System.out.println("relNode:" + relNode);


            System.out.println("----------------------------------------------------------");


            List<String> targetColumnList = tbEnv.from(sinkTbName).getResolvedSchema().getColumnNames();
            List<RelDataTypeField> fieldList = relNode.getRowType().getFieldList();

            System.out.println("targetColumnList:" + targetColumnList);
            System.out.println("sourceColumnList:" + fieldList);

            System.out.println("----------------------------------------------------------");

            RelMetadataQuery metadataQuery = relNode.getCluster().getMetadataQuery();
            List<LineageRel> resultList = new ArrayList<>();


            System.out.println();

            for (int i = 0; i < targetColumnList.size(); i++) {

                String targetColumn = targetColumnList.get(i);

                Set<RelColumnOrigin> columnOrigins = metadataQuery.getColumnOrigins(relNode, i);
                System.out.println(targetColumnList.get(i) + " source:" + columnOrigins.size());

                for (RelColumnOrigin columnOrigin : columnOrigins) {

                    RelOptTable originTable = columnOrigin.getOriginTable();
                    String tableName = Strings.join(".", originTable.getQualifiedName());


                    int ordinal = columnOrigin.getOriginColumnOrdinal();

                    TableSourceTable tableSourceTable = (TableSourceTable) originTable;
                    List<String> columnNames = tableSourceTable.catalogTable().getResolvedSchema().getColumnNames();
                    String sourceColumn = columnNames.get(ordinal);

                    Map<String, String> options = tableSourceTable.catalogTable().getOptions();
                    String sourceConnector = options.get("connector");


                    // add record
                    resultList.add(LineageRel.build(
                            tableName, sourceColumn,
                            sinkTbName, targetColumn,
                            null));

                }
            }


            System.out.println(resultList.size());


            LineageResult lineageResult = buildLineage(resultList);
            JSONObject jobj = (JSONObject)JSONObject.toJSON(lineageResult);
            System.out.println(jobj.toJSONString());

        } else if (!operationType.equals(SqlType.SELECT) && !operationType.equals(SqlType.PRINT)) {
            tbEnv.executeSql(singleSql);

            if (operationType.equals(SqlType.CREATE)) {
                Operation operation = tbEnvImpl.getParser().parse(singleSql).get(0);
                if (operation instanceof CreateTableOperation) {
                    CreateTableOperation createTableOperation = (CreateTableOperation) operation;
                    String connector = createTableOperation.getCatalogTable().getOptions().get("connector");
                    System.out.println("option:" + createTableOperation.getTableIdentifier().asSummaryString() + " " + connector);
                }
            }



        }


    }



    public static SqlType getOperationType(String sql) {
        String sqlTrim = sql.replaceAll(SQL_EMPTY_STR, " ").trim().toUpperCase();
        return Arrays.stream(SqlType.values())
                .filter(sqlType -> sqlType.match(sqlTrim))
                .findFirst()
                .orElse(SqlType.UNKNOWN);
    }




    public static LineageResult buildLineage(List<LineageRel> lineageRelList) {
        List<LineageRelation> relations = new ArrayList<>();
        Map<String, LineageTable> tableMap = new HashMap<>();
        int tableIndex = 1;
        int relIndex = 1;
        for (LineageRel lineageRel : lineageRelList) {
            String sourceTablePath = lineageRel.getSourceTablePath();
            String sourceTableId = null;
            String targetTableId = null;
            if (tableMap.containsKey(sourceTablePath)) {
                LineageTable lineageTable = tableMap.get(sourceTablePath);
                LineageColumn lineageColumn =
                        LineageColumn.build(lineageRel.getSourceColumn(), lineageRel.getSourceColumn());
                if (!lineageTable.getColumns().contains(lineageColumn)) {
                    lineageTable.getColumns().add(lineageColumn);
                }
                sourceTableId = lineageTable.getId();
            } else {
                tableIndex++;
                LineageTable lineageTable = LineageTable.build(String.valueOf(tableIndex), sourceTablePath, "XXX");
                lineageTable
                        .getColumns()
                        .add(LineageColumn.build(lineageRel.getSourceColumn(), lineageRel.getSourceColumn()));
                tableMap.put(sourceTablePath, lineageTable);
                sourceTableId = lineageTable.getId();
            }
            String targetTablePath = lineageRel.getTargetTablePath();
            if (tableMap.containsKey(targetTablePath)) {
                LineageTable lineageTable = tableMap.get(targetTablePath);
                LineageColumn lineageColumn =
                        LineageColumn.build(lineageRel.getTargetColumn(), lineageRel.getTargetColumn());
                if (!lineageTable.getColumns().contains(lineageColumn)) {
                    lineageTable.getColumns().add(lineageColumn);
                }
                targetTableId = lineageTable.getId();
            } else {
                tableIndex++;
                LineageTable lineageTable = LineageTable.build(String.valueOf(tableIndex), targetTablePath, "XXX");
                lineageTable
                        .getColumns()
                        .add(LineageColumn.build(lineageRel.getTargetColumn(), lineageRel.getTargetColumn()));
                tableMap.put(targetTablePath, lineageTable);
                targetTableId = lineageTable.getId();
            }
            LineageRelation lineageRelation = LineageRelation.build(
                    sourceTableId, targetTableId, lineageRel.getSourceColumn(), lineageRel.getTargetColumn());
            if (!relations.contains(lineageRelation)) {
                relIndex++;
                lineageRelation.setId(relIndex + "");
                relations.add(lineageRelation);
            }
        }
        List<LineageTable> tables = new ArrayList<LineageTable>(tableMap.values());
        return LineageResult.build(tables, relations);
    }
}
